package http3

import (
	"time"
)

// CustomCongestionController 实现自定义拥塞控制
type CustomCongestionController struct {
	rtt             time.Duration
	cwnd            uint64
	ssthresh        uint64
	minRTT          time.Duration
	maxDatagramSize uint64
	lossRate        float64
}

func NewCongestionController() *CustomCongestionController {
	return &CustomCongestionController{
		cwnd:            10 * 1460, // 初始窗口大小
		ssthresh:        uint64(1 << 30),
		maxDatagramSize: 1460,
		lossRate:        0.0,
	}
}

func (c *CustomCongestionController) OnPacketSent(sentTime time.Time, bytesInFlight uint64, packetNumber uint64, bytes uint64, isRetransmittable bool) {
	if bytesInFlight+bytes <= c.cwnd {
		return
	}

	if c.lossRate > 0.1 {
		c.cwnd = c.cwnd / 2
	}
}

func (c *CustomCongestionController) OnPacketAcked(number uint64, ackedBytes uint64, rtt time.Duration) {
	if c.rtt == 0 || rtt < c.minRTT {
		c.minRTT = rtt
	}
	c.rtt = rtt

	if c.cwnd < c.ssthresh {
		c.cwnd += c.maxDatagramSize
	} else {
		c.cwnd += (c.maxDatagramSize * c.maxDatagramSize) / c.cwnd
	}
}

func (c *CustomCongestionController) OnPacketLost(number uint64, lostBytes uint64) {
	c.lossRate = (c.lossRate*0.7 + 0.3)
	c.ssthresh = c.cwnd / 2
	c.cwnd = c.ssthresh
}

func (c *CustomCongestionController) SetMaxDatagramSize(size uint64) {
	c.maxDatagramSize = size
}

func (c *CustomCongestionController) GetCongestionWindow() uint64 {
	return c.cwnd
}

func (c *CustomCongestionController) TimeUntilSend(bytesInFlight uint64) time.Time {
	if bytesInFlight >= c.cwnd {
		return time.Now().Add(c.rtt)
	}
	return time.Time{}
}

func (c *CustomCongestionController) HasPacingBudget() bool {
	return true
}

func (c *CustomCongestionController) OnCongestionEvent(number uint64, lostBytes uint64, rtt time.Duration) {
	c.OnPacketLost(number, lostBytes)
}

func (c *CustomCongestionController) InSlowStart() bool {
	return c.cwnd < c.ssthresh
}

func (c *CustomCongestionController) InRecovery() bool {
	return c.lossRate > 0
}
